


#测试用一个图片输入，跳过第一部分，便于答辩展示
import sys

from PIL import Image, ImageDraw,ImageFont 
import torch
import torchvision
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
import torchvision.transforms as tT
from torchvision.models.detection.keypoint_rcnn import KeypointRCNNPredictor
# 导入本文件写的一些函数接口
from _pointer_meter_helpers import my_NMS,remove_low_scores,get_center_seq,pointer_to_read


filePath='00_jepg2.jpg'
img = Image.open(filePath)
# img = rotate_im_accord_exiftag(img)#处理手机拍照时的旋转问题
img = img.convert('RGB')#如果不使用.convert('RGB')进行转换的话，读出来的图像是RGBA四通道的，A通道为透明通道
im_sz=(500,500)
img = img.resize(im_sz)

# # 首先定义一个转换，将 PIL.Image 转换为 PyTorch 张量  
transform = tT.Compose([  
    tT.ToTensor()  # 将 PIL.Image 或 ndarray 转换为 torch.Tensor，并归一化到 [0.0, 1.0]  
])  

# 将 PIL.Image 转换为 PyTorch 张量  
img = transform(img) 
x = img.unsqueeze(0)#.unsqueeze(0)增加维度（0表示，在第一个位置增加维度）




# 2个类别，背景和指针盘
num_classes = 2
# 2个关键点
num_keypoints = 2
# 加载模型预训练的关键点检测模型
model = torchvision.models.detection.keypointrcnn_resnet50_fpn(weights=None) 
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# 加载预训练权重
model.load_state_dict(torch.load('./weights/keypointrcnn_resnet50_fpn_coco-fc266e95.pth', map_location=device)) 
# 获取输入特征数并用新的头替换预先训练好的头
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
in_features2 = model.roi_heads.keypoint_predictor.kps_score_lowres.in_channels
model.roi_heads.keypoint_predictor = KeypointRCNNPredictor(in_features2, num_keypoints)

exp_name = 'pt-dir-detection'
exp_no_2='04'
fn2 = './weights/model_weights_'+exp_name+'_'+exp_no_2+'.pth'

model.load_state_dict(torch.load(fn2, map_location=torch.device(device=device)))
model.to(device)
model.eval()

# # 将 PIL.Image 转换为 PyTorch 张量  

x = x.to(device)
predictions = model(x)
#计算预测值
boxes = predictions[0]['boxes']
scores = predictions[0]['scores']

#置信度得分低于score_threshold的去除
score_threshold=0.9
boxes,scores=remove_low_scores(boxes,scores,score_threshold)
#iou大于这个阈值则去除
nms_threshold = 0.1
selected_idx = my_NMS(boxes, scores, nms_threshold,8)#这个是我自己写的NMS函数，8是要选取的目标个数8个框

keypoints = predictions[0]['keypoints'].cpu().detach().numpy()
boxes = boxes.cpu().detach().numpy()
total=get_center_seq(keypoints,selected_idx)
predict_value=pointer_to_read(total,8)#按照标签读数
predict_read=0
for j in range(8):
    predict_read=predict_read*10+predict_value[j]  

# 将PIL图像转换为Qt可用的QImage
img_pre=tT.ToPILImage()(img.squeeze(0))  # 将 Tensor 转换回 PIL.Image，并移除增加的维度

draw = ImageDraw.Draw(img_pre)

for i in selected_idx:
    # 绘制边界框
    box = boxes[i]
    draw.rectangle(box, outline="red")
    draw.text((10,10),f"predict_read:{predict_read}",font=ImageFont.load_default(),fill="red")
    # 绘制指针线
    draw.line(((keypoints[i,0,0],keypoints[i,0,1]),keypoints[i,1,0],keypoints[i,1,1]), fill="blue")
img_pre.save("debug_image.png")
print('predict_read',predict_read)

